import json
import triton_python_backend_utils as pb_utils
import cv2
import onnxruntime
import numpy as np
import os
import logging
from logging.handlers import RotatingFileHandler

log_filename = '/model.log'
onnxruntime.set_default_logger_severity(3)

file_handler = RotatingFileHandler(log_filename, maxBytes=50 * 1024 * 1024, backupCount=5)
console_handler = logging.StreamHandler()

log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(log_formatter)
console_handler.setFormatter(log_formatter)

logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(file_handler)
logger.addHandler(console_handler)


def add_log_separator():
    logger.info("-" * 80)

local_path = os.path.dirname(os.path.abspath(__file__))
model_filename = "model.onnx"
onnx_model_path = os.path.join(local_path, model_filename)

class TritonPythonModel:
    def initialize(self, args):
        try:
            if args is None:
                raise ValueError("Received 'None' for 'args'")
            if 'model_config' not in args:
                raise ValueError("Expected 'model_config' in 'args', but not found")

            self.model_config = json.loads(args['model_config'])
            out_bboxes_config = pb_utils.get_output_config_by_name(self.model_config, "bboxes")
            out_scores_config = pb_utils.get_output_config_by_name(self.model_config, "scores")
            out_classes_config = pb_utils.get_output_config_by_name(self.model_config, "classes")
            out_labels_config = pb_utils.get_output_config_by_name(self.model_config, "labels")

            self.out_bboxes_dtype = pb_utils.triton_string_to_numpy(out_bboxes_config['data_type'])
            self.out_scores_dtype = pb_utils.triton_string_to_numpy(out_scores_config['data_type'])
            self.out_classes_dtype = pb_utils.triton_string_to_numpy(out_classes_config['data_type'])
            self.out_labels_dtype = pb_utils.triton_string_to_numpy(out_labels_config['data_type'])
            logger.info(f">> The model is loading...")
            # self.sess = onnxruntime.InferenceSession(onnx_model_path)
            self.sess = onnxruntime.InferenceSession(onnx_model_path, providers=['CUDAExecutionProvider'])
            logger.info(f">> Model loaded successfully!")
        except Exception as e:
            logger.error(f'>> Failed to initialize model: {e}')
            raise

    def execute(self, requests):
        responses = []
        for request in requests:
            try:
                add_log_separator()
                logger.info(f">> Received request...")
                image = pb_utils.get_input_tensor_by_name(request, 'image').as_numpy()
                score = pb_utils.get_input_tensor_by_name(request, 'score')
                if score is None:
                    score = np.float32([0.3])
                    logger.info(">> No score input, use default value 0.3")
                else:
                    score = score.as_numpy().astype(np.float32)
                logger.info(f">> Input image shape: {image.shape}, score: {score}")
                result, img_after = self.Inference(image, score)
                # whether the result is empty
                if result[0][0].size == 0:
                    logger.info(">> No object detected!!!")
                    e_bboxes = np.empty((0, 4), dtype=self.out_bboxes_dtype)
                    e_conf = np.empty((0, ), dtype=self.out_scores_dtype)
                    e_cls = np.empty((0, ), dtype=self.out_classes_dtype)
                    e_labels = np.empty((0, 0), dtype=self.out_labels_dtype)
                    out_tensor_bboxes = pb_utils.Tensor('bboxes', e_bboxes)
                    out_tensor_scores = pb_utils.Tensor('scores', e_conf)
                    out_tensor_classes = pb_utils.Tensor('classes', e_cls)
                    out_tensor_labels = pb_utils.Tensor('labels', e_labels)
                    inference_response = pb_utils.InferenceResponse(
                        output_tensors=[out_tensor_bboxes, out_tensor_scores, out_tensor_classes, out_tensor_labels])
                    responses.append(inference_response)
                    continue
                bboxes = self.cod_trf(result[0], image, img_after)
                conf = np.array(result[1][0] * 100)
                cls = np.array(result[2][0])
                bbox = np.array(np.round(bboxes).astype(np.uint32))
                logger.info(f">> The number of detected targets: {len(bbox)}")
                # logger.info(f">> Inference results: bboxes:{bbox} conf:{conf} cls:{cls}")
                labels_name = [
                    "人员",
                    "自行车",
                    "汽车",
                    "摩托车",
                    "飞机",
                    "公交车",
                    "火车",
                    "卡车",
                    "船",
                    "交通灯",
                    "消防栓",
                    "停车标志",
                    "停车计时器",
                    "长椅",
                    "鸟",
                    "猫",
                    "狗",
                    "马",
                    "羊",
                    "牛",
                    "大象",
                    "熊",
                    "斑马",
                    "长颈鹿",
                    "背包",
                    "雨伞",
                    "手提包",
                    "领带",
                    "行李箱",
                    "飞盘",
                    "滑雪板",
                    "滑雪板",
                    "运动球",
                    "风筝",
                    "棒球棒",
                    "棒球手套",
                    "滑板",
                    "冲浪板",
                    "网球拍",
                    "瓶子",
                    "酒杯",
                    "杯子",
                    "叉子",
                    "刀",
                    "勺子",
                    "碗",
                    "香蕉",
                    "苹果",
                    "三明治",
                    "橙子",
                    "西兰花",
                    "胡萝卜",
                    "热狗",
                    "披萨",
                    "甜甜圈",
                    "蛋糕",
                    "椅子",
                    "沙发",
                    "盆栽",
                    "床",
                    "餐桌",
                    "马桶",
                    "电视",
                    "笔记本电脑",
                    "鼠标",
                    "遥控器",
                    "键盘",
                    "手机",
                    "微波炉",
                    "烤箱",
                    "烤面包机",
                    "水槽",
                    "冰箱",
                    "书",
                    "时钟",
                    "花瓶",
                    "剪刀",
                    "泰迪熊",
                    "吹风机",
                    "牙刷",
                ]
                labels = [labels_name[c] for c in cls]
                labels = np.array(labels, dtype=object)
                # 过滤指定类别[1,2,3,5,7]
                out_bboxes = []
                out_label = []
                out_classes = []
                out_score = []
                filter_index = [1, 2, 3, 5, 7]
                for i, clss in enumerate(cls):
                    if clss in filter_index:
                        out_bboxes.append(bbox[i])
                        out_label.append(labels[i])
                        out_classes.append(clss)
                        out_score.append(conf[i])
                out_tensor_bboxes = pb_utils.Tensor('bboxes', np.array(out_bboxes).astype(self.out_bboxes_dtype))
                out_tensor_scores = pb_utils.Tensor('scores', np.array(out_score).astype(self.out_scores_dtype))
                out_tensor_classes = pb_utils.Tensor('classes', np.array(out_classes).astype(self.out_classes_dtype))
                out_tensor_labels = pb_utils.Tensor('labels', np.array(out_label).astype(self.out_labels_dtype))
                inference_response = pb_utils.InferenceResponse(
                    output_tensors=[out_tensor_bboxes, out_tensor_scores, out_tensor_classes, out_tensor_labels])
                responses.append(inference_response)
                logger.info(f">> Response executed successfully!")
            except Exception as e:
                logger.error(f'>> Failed to execute request: {e}')
                inference_response = pb_utils.InferenceResponse(output_tensors=[], error=pb_utils.TritonError(str(e)))
                responses.append(inference_response)
        return responses

    def finalize(self):
        logger.info('>> Cleaning up...')

    def Inference(self, image, score):
        try:
            std_h, std_w = 640, 640
            img_after = self.resize_image(image, (std_w, std_h), True)  # （640， 640， 3）
            imageData = self.img2input(img_after)

            result = self.sess.run(None, {'images': imageData, 'score': score})
            return result, img_after

        except Exception as e:
            logger.error(f">> Inference failed: {e}")
            raise

    def resize_image(self, image, size, letterbox_image):
        ih, iw, _ = image.shape
        h, w = size
        if letterbox_image:
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)
            image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_LINEAR)
            image_back = np.ones((h, w, 3), dtype=np.uint8) * 128
            image_back[(h - nh) // 2: (h - nh) // 2 + nh, (w - nw) // 2:(w - nw) // 2 + nw, :] = image
        else:
            image_back = image
        return image_back

    def img2input(self, img):
        img = np.transpose(img, (2, 0, 1))
        img = img / 255
        return np.expand_dims(img, axis=0).astype(np.float32)

    def xywh2xyxy(self, *box):
        ret = [box[0] - box[2] // 2, box[1] - box[3] // 2, box[0] + box[2] // 2, box[1] + box[3] // 2]
        return ret

    def cod_trf(self, result, pre, after):
        result_boxes = []
        for i in result[0]:
            x, y, w, h = i
            x1, y1, x2, y2 = self.xywh2xyxy(x, y, w, h)
            h_pre, w_pre, _ = pre.shape
            h_after, w_after, _ = after.shape
            scale = max(w_pre / w_after, h_pre / h_after)
            h_pre, w_pre = h_pre / scale, w_pre / scale
            x_move, y_move = abs(w_pre - w_after) // 2, abs(h_pre - h_after) // 2
            ret_x1, ret_x2 = (x1 - x_move) * scale, (x2 - x_move) * scale
            ret_y1, ret_y2 = (y1 - y_move) * scale, (y2 - y_move) * scale
            ret = np.array([ret_x1, ret_y1, ret_x2, ret_y2])
            result_boxes.append(ret)
        result[0] = result_boxes
        return result_boxes


